// https://leetcode.cn/problems/kth-largest-element-in-an-array/description/

// 算法思路总结：
// 1. 基于三路分区的快速选择算法寻找第K大元素
// 2. 随机选择基准值key，将数组分为三部分：<key, =key, >key
// 3. 根据K与三部分大小的关系，决定递归搜索的区域
// 4. 第K大元素在降序排列中的位置，即第(nums.size()-K+1)小
// 5. 时间复杂度：O(n)，空间复杂度：O(logn)

#include <iostream>
using namespace std;

#include <vector>
#include <algorithm>
#include <ctime>
#include <cstdlib>

class Solution 
{
public:
    int findKthLargest(vector<int>& nums, int k) 
    {
        srand((unsigned)time(NULL));
        return qselect(nums, 0, nums.size() - 1, k);
    }

    int getRandom(vector<int>& nums, int left, int right)
    {
        int r = rand();
        return nums[r % (right - left + 1) + left];
    }

    int qselect(vector<int>& nums, int left, int right, int k)
    {
        if (left == right)
        {
            return nums[left];
        }

        int key = getRandom(nums, left, right);
        int begin = left - 1, end = right + 1, cur = left;
        while (cur < end)
        {
            if (nums[cur] < key)
            {
                swap(nums[++begin], nums[cur++]);
            }
            else if (nums[cur] == key)
            {
                cur++;
            }
            else if (nums[cur] > key)
            {
                swap(nums[--end], nums[cur]);
            }
        }

        int a = begin - left + 1, b = end - begin - 1, c = right - end + 1;
        if (k <= c) return qselect(nums, end, right, k);
        else if (k <= b + c) return key;
        else return qselect(nums, left, begin, k - b - c);
    }   
};

int main()
{
    vector<int> nums1 = {3,2,1,5,6,4}, nums2 = {3,2,3,1,2,4,5,5,6};
    int k1 = 2, k2 = 4;

    Solution sol;

    cout << sol.findKthLargest(nums1, k1) << endl;
    cout << sol.findKthLargest(nums2, k2) << endl;

    return 0;
}